package hr.irb.goTermList;

import hr.irb.geneOntology.GOTerm;
import hr.irb.geneOntology.GeneOntology;
import java.io.BufferedReader;
import java.io.File;
import java.io.FileInputStream;
import java.io.FileNotFoundException;
import java.io.FileOutputStream;
import java.io.FileReader;
import java.io.IOException;
import java.io.InputStream;
import java.io.InputStreamReader;
import java.io.ObjectInput;
import java.io.ObjectInputStream;
import java.io.ObjectOutput;
import java.io.ObjectOutputStream;
import java.io.OutputStream;
import java.io.Serializable;
import java.util.Collection;
import java.util.HashMap;
import java.util.HashSet;
import java.util.List;
import java.util.Map;
import java.util.Set;
import java.util.zip.GZIPInputStream;
import java.util.zip.GZIPOutputStream;

/**
 *
 * @author Fran Supek
 */
public class GoTermSizes implements Serializable {

  private static final long serialVersionUID = 4L;

  /** A Map, mapping from GO id to size (# of proteins annotated with it). */
  protected Map<Integer, Integer> termSizes = new HashMap<Integer, Integer>();

  /**
   * Category sizes, divided by the size of the root node in that particular
   * namespace.
   */
  protected Map<Integer, Double> termSizesNormalized = null;

  public final int sumOfSizes;
  public final double sumOfNormalizedSizes;

  public boolean isSizeKnown(int goId) {
    return termSizes.containsKey(goId);
  }

  public int getSize(int goId) {
    if ( termSizes.containsKey(goId) ) {
      return termSizes.get(goId);
    } else {
      //return sumOfSizes / termSizes.size();  // return average size, if no data given
      throw new IllegalArgumentException("GO category "+goId+" does not have a defined size.");
    }
  }


  public int getSizeGuessIfUnknown(int goId, GeneOntology myGo) {
    if ( termSizes.containsKey(goId) ) {
      return termSizes.get(goId);
    } else {
      // look at sizes of all siblings, and average them to guess the
      // frequency of the term
      GOTerm curTerm = myGo.get(goId);
      Set<GOTerm> sibs = curTerm.getSiblings();
      int numUsableSibs = 0;
      int sumOfSizesOfUsableSibs = 0;
      for ( GOTerm sib : sibs ) {
        if ( termSizes.containsKey(sib.getId())
                && termSizes.get(sib.getId()) > 0 ) {
          numUsableSibs++;
          sumOfSizesOfUsableSibs += termSizes.get(sib.getId());
        }
      }
      if ( numUsableSibs >= 1 ) {
        this.termSizes.put(goId, sumOfSizesOfUsableSibs/numUsableSibs);
        return sumOfSizesOfUsableSibs / numUsableSibs;
      }
      // no usable siblings? try frequency of largest child
      Set<GOTerm> children = curTerm.getChildren();
      int sizeOfBiggestChild = 0;
      for ( GOTerm child : children ) {
        if ( termSizes.containsKey(child.getId()) ) {
          if ( termSizes.get(child.getId()) > sizeOfBiggestChild )
            sizeOfBiggestChild = termSizes.get(child.getId());
        }
      }
      if (sizeOfBiggestChild > 0) {
        this.termSizes.put(goId, sizeOfBiggestChild);
        return sizeOfBiggestChild;
      }
      // no usable children? try frequency of smallest parent
      Set<GOTerm> parents = curTerm.getAllParents();
      int sizeOfSmallestParent = Integer.MAX_VALUE;
      for ( GOTerm parent : parents ) {
        if ( termSizes.containsKey(parent.getId()) ) {
          if ( termSizes.get(parent.getId()) > 0
                  && termSizes.get(parent.getId()) < sizeOfSmallestParent )
            sizeOfSmallestParent = termSizes.get(parent.getId());
        }
      }
      if (sizeOfSmallestParent < Integer.MAX_VALUE) {
        this.termSizes.put( goId, sizeOfSmallestParent );
        return sizeOfSmallestParent;
      }
      // No more options to try... this happens in very rare cases with terms
      // that have been made obsolete before being officially introduced
      // and that don't have any parents or children assigned. In that case,
      // return the average size of all GO Terms.
      return sumOfSizes / termSizes.size();
      //throw new GeneOntologyException("GO category "+goId+" does not have " +
      //        "a defined size and the size was also impossible to guess.");
    }
  }
  

  public double getFreq(int goId) {
    if ( termSizesNormalized.containsKey(goId) ) {
      return termSizesNormalized.get(goId);
    } else {
      throw new IllegalArgumentException("GO category "+goId+" does not have a defined frequency.");
    }
  }

  public double getFreqGuessIfUnknown(int goId, GeneOntology myGo) {
    if ( termSizesNormalized.containsKey(goId) ) {
      return termSizesNormalized.get(goId);
    } else {
      // look at frequencies of all siblings, and average them to guess the 
      // frequency of the term
      GOTerm curTerm = myGo.get(goId);
      Set<GOTerm> sibs = curTerm.getSiblings();
      int numUsableSibs = 0;
      double sumOfFreqsOfUsableSibs = 0.0;
      for ( GOTerm sib : sibs ) {
        if ( termSizesNormalized.containsKey(sib.getId())
                && termSizesNormalized.get(sib.getId()) > 0
                && !Double.isNaN(termSizesNormalized.get(sib.getId())) ) {
          numUsableSibs++;
          sumOfFreqsOfUsableSibs += termSizesNormalized.get(sib.getId());
        }
      }
      if ( numUsableSibs >= 1 ) {
        this.termSizesNormalized.put( goId, sumOfFreqsOfUsableSibs / numUsableSibs );
        return sumOfFreqsOfUsableSibs / numUsableSibs;
      }
      // no usable siblings? try frequency of largest child
      Set<GOTerm> children = curTerm.getChildren();
      double freqOfBiggestChild = 0.0;
      for ( GOTerm child : children ) {
        if ( termSizesNormalized.containsKey(child.getId())
                && !Double.isNaN(termSizesNormalized.get(child.getId())) ) {
          if ( termSizesNormalized.get(child.getId()) > freqOfBiggestChild )
            freqOfBiggestChild = termSizesNormalized.get(child.getId());
        }
      }
      if (freqOfBiggestChild > 0.0) {
        this.termSizesNormalized.put( goId , freqOfBiggestChild );
        return freqOfBiggestChild;
      }
      // no usable children? try frequency of smallest parent
      Set<GOTerm> parents = curTerm.getAllParents();
      double freqOfSmallestParent = Double.MAX_VALUE;
      for ( GOTerm parent : parents ) {
        if ( termSizesNormalized.containsKey(parent.getId())
                && !Double.isNaN(termSizesNormalized.get(parent.getId())) ) {
          if ( termSizesNormalized.get(parent.getId()) < freqOfSmallestParent )
            freqOfSmallestParent = termSizesNormalized.get(parent.getId());
        }
      }
      if (freqOfSmallestParent < Double.MAX_VALUE) {
        this.termSizesNormalized.put( goId, freqOfSmallestParent );
        return freqOfSmallestParent;
      }
      // No more options to try... this happens in very rare cases with terms
      // that have been made obsolete before being officially introduced
      // and that don't have any parents or children assigned. In that case,
      // return the average frequency of all GO Terms

      return sumOfNormalizedSizes / termSizesNormalized.size();

      //throw new GeneOntologyException("GO category "+goId+" does not have " +
      //        "a defined frequency and the frequency was also impossible to guess.");
    }

  }
  

  public final GoTermWordCorpus wordCorpus;


  /**
   * Loads the Gene Ontology category sizes by parsing an EBI GO Annotation
   * file (GOA). These files are quite large so this will be slow. You may
   * limit this loading to a set of species by their taxId.<p>
   *
   * Use serialization/deserialization to instantiate these objects much
   * faster.<p>
   *
   * @param goaUniprotGz the file "gene_association.goa_uniprot.gz"
   * @param allowedTaxIds A list of taxIds to limit the counting of
   * annotations; if you pass null or an empty list, all annotations will be
   * loaded (no limits).
   * @param limitToSwissProt CURRENTLY IGNORED! Previous function: Limits
   * results only to high quality Swiss-Prot subsets; otherwise, uses whole of
   * UniProt (Swiss-Prot + TREmbl)
   */
  public GoTermSizes( File goaUniprotGz, Set<Integer> allowedTaxIds,
          boolean limitToSwissProt, GeneOntology myGo )
          throws IOException {

    BufferedReader br;
    if ( goaUniprotGz.getName().toLowerCase().endsWith(".gz") ) {
      GZIPInputStream gzip = new GZIPInputStream(
              new FileInputStream(goaUniprotGz) );
      br = new BufferedReader(new InputStreamReader(gzip));
    } else {
      br = new BufferedReader(new FileReader(goaUniprotGz));
    }

    Set<GOTerm> allGosForCurrentGene = new HashSet<GOTerm>();

    String line = null, lastId = null, lastGo = null, lastTaxon = null;

    while ( (line = br.readLine()) != null ) {

      if ( !line.startsWith("UniProtKB") )
        continue;

      // cols[1] is the protein ID
      // cols[4] is the GO
      // cols[12] is the taxon
      String[] cols = line.split("\t");

      /*if ( limitToSwissProt && cols[0].equals("UniProtKB/TrEMBL") )
        break; // no need to look further, only TrEMBL annotations ahead
      if ( limitToSwissProt && !cols[0].equals("UniProtKB/Swiss-Prot") )
        continue;*/

      // sometimes lines with annotation repeat to account for different
      // sources of annotation (we don't care about the source here)
      if ( cols[1].equals(lastId)
              && cols[4].equals(lastGo) )
        continue;

      if ( !cols[1].equals(lastId) )
        allGosForCurrentGene.clear();

      lastId = cols[1];
      lastGo = cols[4];
      //lastTaxon = cols[12];

      int goIdNumeric;
      try {
        goIdNumeric = Integer.parseInt(cols[4].substring(3));
      } catch (NumberFormatException e) {
        System.err.println("STRANGE ROW:" + line);
        continue; // weird row...
      }

      if ( allowedTaxIds == null || allowedTaxIds.isEmpty() ) {
        increaseCountOfGoCat(goIdNumeric, myGo, allGosForCurrentGene);
        continue;
      }

      int taxonNumeric;
      try {
        taxonNumeric = Integer.parseInt(cols[12].substring(6));
      } catch (NumberFormatException numberFormatException) {
        System.err.println("STRANGE ROW:" + line);
        continue; // weird row...
      }

      if ( allowedTaxIds.contains(taxonNumeric) ) {
        increaseCountOfGoCat(goIdNumeric, myGo, allGosForCurrentGene);
      }


    }

    int sumOfSizesTemp = 0;
    for ( Integer curSize : termSizes.values() ) {
      sumOfSizesTemp += curSize;
    }
    sumOfSizes = sumOfSizesTemp;

    termSizesNormalized = normalizeCategorySizesToRootNode(termSizes, myGo);
    double sumOfSizesNormalizedTemp = 0;
    for ( Double curSize : termSizesNormalized.values() ) {
      sumOfSizesNormalizedTemp += curSize;
    }
    sumOfNormalizedSizes = sumOfSizesNormalizedTemp;

    wordCorpus = new GoTermWordCorpus(this, myGo);

  }


  /**
   * Increases by 1 the count of the GO term with the given GO ID, and
   * also all of its parents. Will not count terms that have already been
   * counted - if they have, for example, been the parents of the currently
   * provided term.
   *
   * @param goId
   * @param myGo
   * @param alreadyAssignedGoTerms
   */
  private void increaseCountOfGoCat(int goId, GeneOntology myGo,
          Set<GOTerm> alreadyAssignedGoTerms){

    GOTerm curGoTerm = myGo.get(goId);
    Set<GOTerm> workingSet = new HashSet<GOTerm>();
    if ( curGoTerm != null ) {
      workingSet.add(curGoTerm);
      Collection<GOTerm> allParents = curGoTerm.getAllParents();
      if ( allParents != null)
        workingSet.addAll(allParents);
    }

    for ( GOTerm aTerm : workingSet ) {
      if ( alreadyAssignedGoTerms.contains(aTerm) )
        continue;
      int curGoId = aTerm.getId();
      if ( termSizes.containsKey(curGoId) ) {
        termSizes.put(curGoId, termSizes.get(curGoId)+1);
      } else {
        termSizes.put(curGoId, 1);
      }
      alreadyAssignedGoTerms.add(aTerm);
    }

  }

  /**
   * Deserializes a GoTermSizes object from a file within the Revigo JAR. 
   * The files are stored in hr.irb.goTermList.goTermSizesData
   * 
   * @param goaPackageName Can currently be "goa_2012_04_02" or (newest) 
   * "goa_2013_04_30".
   * @param taxId A NCBI Taxonomy ID for the organism to load; pass 0 for the
   * entire Uniprot database. (note: E. coli is 83333)
   * @return A deserialized GoTermSizes object.
   * @throws IOException
   * @throws ClassNotFoundException 
   */
  public static GoTermSizes deserializeFromResources ( String goaPackageName, Integer taxId )
          throws IOException, ClassNotFoundException {
      return deserializeFromStream(
              GoTermSizes.class.getResourceAsStream( "goTermSizesData/" + goaPackageName + "/" + taxId + ".ser.gz" ),
              true );  // we assume gzipped ".ser" files    
   }


  
  
  /** Assumes file iz gzipped, iff the filename ends with .gz. */
  public static GoTermSizes deserializeFromFile ( File file )
          throws IOException, ClassNotFoundException {
    InputStream is;
    if ( file.getName().endsWith(".gz") ) {
      is = new GZIPInputStream(new FileInputStream(file));
    } else {
      is = new FileInputStream(file);
    }
    return deserializeFromStream(is, false);
  }

  
  /** The serialized object may be gzipped. */
  public static GoTermSizes deserializeFromStream ( InputStream is, boolean gzipped )
          throws IOException, ClassNotFoundException {

    ObjectInput oi ;
    if ( gzipped ) {
      oi = new ObjectInputStream(new GZIPInputStream(is));
    } else {
      oi = new ObjectInputStream(is);
    }
      
    Object newObj = oi.readObject();
    if ( ! (newObj instanceof GoTermSizes) )
      throw new IllegalArgumentException("Wrong class in file;");
    oi.close();
    return (GoTermSizes)newObj;

  }


  /** Will compress the file if the name ends with ".gz". */
  public void serializeToFile( File file )
          throws FileNotFoundException, IOException {

    ObjectOutput oo;
    if ( file.getName().endsWith(".gz") ) {
      oo = new ObjectOutputStream(
              new GZIPOutputStream(new FileOutputStream(file)));
    } else {
      oo = new ObjectOutputStream(new FileOutputStream(file));
    }
    oo.writeObject(this);
    oo.close();

  }


  /**
   * Makes serialized files with GoTermSizes for a number of model organisms
   * used in REVIGO. Starts from the "gene_association.goa_uniprot.gz" file
   * from the Uniprot GOA project.
   *
   * @param args not used
   */
  public static void main(String[] args) throws IOException {

    int[] taxIds = new int[]{
      83333, // E. coli
      1423,  // B. subtilis
      2097,  // M. genitalium
      155892, // Caulobacter crescentus
      1148,  // Synechocystis sp.
      287,   // P. aeruginosa
      2104,  // Mycoplasma pneumoniae

      9606,  // human
      10116, // rat
      10090, // mouse

      7955,  // Danio rerio (zebrafish)

      7227,  // fruit fly
      6239,  // worm (C. elegans)
      3702,  // arabidopsis
      39947, // Oryza sativa
      4577,  // Zea mays
      4932,  // S. cerevisiae (bakers yeast)
      4896,  // S. pombe (fission yeast)
      9913,   // Bos taurus
      5833,   // Plasmodium falciparum
      31033,  // Takifugu rubripes
      8355,   // Xenopus laevis
      3055,   // Chlamydomonas reinhardtii
      44689   // Dictyostelium discoideum*/
    };

    GeneOntology myGo = new GeneOntology();
    
    myGo.addBasicKeywordsToAllTerms();
    myGo.addKeywordsFromUniprotKeywords( );
    
    // loop through taxIds of interest, and serialize one-by-one
    for ( int curTaxId : taxIds ) {

      Set<Integer> taxIdsParam = new HashSet<Integer>();
      taxIdsParam.add(curTaxId);

      GoTermSizes sizes = new GoTermSizes(
              //new File("Z:/databases/EBI_GOA_mappings/goa_2012-04-02/" +
              new File("Z:/databases/EBI_GOA_mappings/goa_2013-04-30/" +
              "gene_association.goa_uniprot.gz"),
              taxIdsParam, false, myGo);

      sizes.serializeToFile(
              //new File("Z:/databases/EBI_GOA_mappings/goa_2012-04-02/goTermSizes/"
              new File("Z:/databases/EBI_GOA_mappings/goa_2013-04-30/goTermSizes/"
              + curTaxId + ".ser.gz")
              );

      System.out.println("Done with " + curTaxId);

    }

    GoTermSizes sizes = new GoTermSizes(
            //new File("Z:/databases/EBI_GOA_mappings/goa_2012-04-02/" +
            new File("Z:/databases/EBI_GOA_mappings/goa_2013-04-30/" +
            "gene_association.goa_uniprot.gz"),
            null, false, myGo);
    
    System.out.printf("%s\t%s\t%s", "goId", "IC", "desc" );
    for ( Integer goId : sizes.termSizesNormalized.keySet() ) {
      System.out.printf("%d\t%f\t%s\n", goId, -Math.log(sizes.termSizesNormalized.get(goId))/Math.log(2), myGo.get(goId) );
    }
      
    // also serialize the "whole uniprot" file
    sizes.serializeToFile(
            //new File("Z:/databases/EBI_GOA_mappings/goa_2012-04-02/goTermSizes/"
            new File("Z:/databases/EBI_GOA_mappings/goa_2013-04-30/goTermSizes/"
            + "0.ser.gz")
            );
     
  }



  private Map<Integer, Double> normalizeCategorySizesToRootNode(
          Map<Integer, Integer> catSizes, GeneOntology myGo) {

    Map<Integer, Double> result = new HashMap<Integer, Double>();

    for ( int curGoId : catSizes.keySet() ) {
      
      GOTerm curGo = myGo.get(curGoId);
      if ( curGo == null )
        continue;
      
      GOTerm rootOfCurGo = (GOTerm) curGo.getTopNode();
      double normSize = (double) catSizes.get(curGoId) / (double) catSizes.get(rootOfCurGo.getId());
      result.put( curGoId, normSize );

    }

    return result;

  }


}
